我們今天換來SensorFlow來實現一個簡單的超分辨率卷積神經網絡(SRCNN) 模型!
import tensorflow as tf
from tensorflow.keras import layers, models
import numpy as np
from PIL import Image
定義 SRCNN 模型
def build_srcnn_model():
model = models.Sequential()
model.add(layers.Conv2D(64, (9, 9), activation='relu', padding='same', input_shape=(None, None, 1)))
model.add(layers.Conv2D(32, (5, 5), activation='relu', padding='same'))
model.add(layers.Conv2D(1, (5, 5), padding='same'))
model.compile(optimizer='adam', loss='mean_squared_error')
return model
載入圖像並轉換為張量
def load_image(img_path):
img = Image.open(img_path).convert('YCbCr')
y, _, _ = img.split()
y = np.array(y).astype(np.float32) / 255.0
return np.expand_dims(np.expand_dims(y, axis=0), axis=-1)
保存處理後的圖像
def save_image(tensor, path):
img = tensor[0, :, :, 0]
img = (img * 255.0).clip(0, 255).astype(np.uint8)
img = Image.fromarray(img)
img.save(path)
初始化模型與損失函數
model = build_srcnn_model()
載入訓練數據
input_image = load_image('low_res_image.jpg')
target_image = load_image('high_res_image.jpg')
訓練模型
num_epochs = 100
model.fit(input_image, target_image, epochs=num_epochs, verbose=1)
生成結果圖像
output_image = model.predict(input_image)
save_image(output_image, 'output_image.jpg')